import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset
import pickle

# 1. 正确读取CSV文件（使用逗号分隔）
df = pd.read_csv('faults.csv', sep=',')

# 2. 打印实际列名以确认
print("实际列名:", df.columns.tolist())

# 3. 分离特征和标签 - 使用正确的列名列表
feature_columns = [
    'X_Minimum', 'X_Maximum', 'Y_Minimum', 'Y_Maximum', 'Pixels_Areas',
    'X_Perimeter', 'Y_Perimeter', 'Sum_of_Luminosity', 'Minimum_of_Luminosity',
    'Maximum_of_Luminosity', 'Length_of_Conveyer', 'TypeOfSteel_A300',
    'TypeOfSteel_A400', 'Steel_Plate_Thickness', 'Edges_Index', 'Empty_Index',
    'Square_Index', 'Outside_X_Index', 'Edges_X_Index', 'Edges_Y_Index',
    'Outside_Global_Index', 'LogOfAreas', 'Log_X_Index', 'Log_Y_Index',
    'Orientation_Index', 'Luminosity_Index', 'SigmoidOfAreas'
]

label_columns = [
    'Pastry', 'Z_Scratch', 'K_Scatch', 'Stains',
    'Dirtiness', 'Bumps', 'Other_Faults'
]

# 4. 提取特征和标签
features = df[feature_columns].values.astype(np.float32)
labels = df[label_columns].values.astype(np.float32)

# 5. 转换为PyTorch张量
features_tensor = torch.tensor(features)
labels_tensor = torch.tensor(labels)

# 8. 保存为PKL文件（只保存张量数据，而不是Dataset对象）
with open('steel_defects_data.pkl', 'wb') as f:
    pickle.dump({
        'features': features_tensor,
        'labels': labels_tensor
    }, f)

print("数据已成功保存为 steel_defects_data.pkl")
print(f"特征形状: {features_tensor.shape}")
print(f"标签形状: {labels_tensor.shape}")

